from typing import Dict, Any, Tuple
from threading import Lock

from axiom_boot.di import service, autowired
from axiom_boot.vector.parser.node_parser_interface import NodeParserInterface
from axiom_boot.vector.parser.parser_type import ParserType
from axiom_boot.vector.parser.sentence_splitter.sentence_splitter_config import SentenceSplitterConfig
from axiom_boot.vector.parser.sentence_splitter.sentence_splitter_parser import SentenceSplitterParser
from axiom_boot.vector.parser.markdown.markdown_node_parser import MarkdownNodeParser
from axiom_boot.vector.parser.markdown.markdown_node_parser_config import MarkdownNodeParserConfig
from axiom_boot.vector.parser.semantic_splitter.semantic_splitter_config import SemanticSplitterConfig
from axiom_boot.vector.parser.semantic_splitter.semantic_splitter_parser import SemanticSplitterParser
from axiom_boot.vector.parser.token_splitter.token_text_splitter_config import TokenTextSplitterConfig
from axiom_boot.vector.parser.token_splitter.token_text_splitter_parser import TokenTextSplitterParser
from axiom_boot.vector.embedding.embedding_manager import EmbeddingManager

ParserCacheKey = Tuple[ParserType, Any]

@service()
class NodeParserManager:
    """
    一个管理器服务，负责按需创建和缓存多种类型的 NodeParser 实例。
    它现在是一个能够根据指定的 ParserType 和可选配置来提供解析器的工厂。
    """

    embedding_manager: EmbeddingManager = autowired()

    def __init__(self):
        self._parsers: Dict[ParserCacheKey, NodeParserInterface] = {}
        self._lock = Lock()

    def get_parser(
        self,
        parser_type: ParserType,
        config: Any | None = None,
    ) -> NodeParserInterface:
        """
        根据解析器类型和可选配置，获取一个 NodeParser 实例。
        """
        # [关键修复] 确保缓存键是可哈希的
        # 如果 config 是字典，将其转换为 frozenset of items，因为 dict 是不可哈希的
        if isinstance(config, dict):
            cache_key_config = frozenset(config.items())
        else:
            cache_key_config = config
        
        # 使用 (类型, 不可变配置) 作为缓存键
        cache_key = (parser_type, cache_key_config)

        if cache_key not in self._parsers:
            with self._lock:
                if cache_key not in self._parsers:
                    self._parsers[cache_key] = self._create_parser(
                        parser_type, config
                    )
        
        return self._parsers[cache_key]

    def _create_parser(
        self,
        parser_type: ParserType,
        config: Any | None = None,
    ) -> NodeParserInterface:
        """
        根据类型和配置创建一个新的解析器实例。
        """
        config_dict = config or {}

        # [统一处理] 在创建任何解析器之前，对分隔符进行反转义
        if 'separator' in config_dict and isinstance(config_dict['separator'], str):
            config_dict['separator'] = config_dict['separator'].encode().decode('unicode_escape')
        
        if 'paragraph_separator' in config_dict and isinstance(config_dict['paragraph_separator'], str):
            config_dict['paragraph_separator'] = config_dict['paragraph_separator'].encode().decode('unicode_escape')

        if parser_type == ParserType.SENTENCE_SPLITTER:
            # [关键修复] 无论传入的是 dict 还是 Pydantic 模型，都用模型类重新构造，确保类型正确
            final_config = SentenceSplitterConfig(**config_dict) if isinstance(config_dict, dict) else (config or SentenceSplitterConfig())
            return SentenceSplitterParser(final_config)
        
        if parser_type == ParserType.TOKEN_SPLITTER:
            final_config = TokenTextSplitterConfig(**config_dict) if isinstance(config_dict, dict) else (config or TokenTextSplitterConfig())
            return TokenTextSplitterParser(final_config)
        
        if parser_type == ParserType.MARKDOWN_NODE:
            final_config = MarkdownNodeParserConfig(**config_dict) if isinstance(config_dict, dict) else (config or MarkdownNodeParserConfig())
            return MarkdownNodeParser(final_config)

        if parser_type == ParserType.SEMANTIC_SPLITTER:
            # [关键修复] 确保我们总是使用 Pydantic 模型实例
            config_obj = SemanticSplitterConfig(**config_dict) if isinstance(config_dict, dict) else (config or SemanticSplitterConfig())
            
            if not config_obj.embedding_profile_name:
                raise ValueError("使用 semantic_splitter 时，必须在配置中指定 embedding_profile_name。")
            
            embedding_service = self.embedding_manager.get_embedding(config_obj.embedding_profile_name)
            
            return SemanticSplitterParser(
                config=config_obj,
                embedding_service=embedding_service
            )
            
        raise TypeError(f"不支持的节点解析器类型: {parser_type}")